import asyncio
from json_utils import try_parse_json_object
from tqdm.auto import tqdm


async def agenerate(client, user_content, return_json=True):
    # 发起异步请求
    response = client.chat.asyncCompletions.create(
        model="glm-4-plus",  # 指定模型名称
        messages=[
            {
                "role": "user",
                "content": user_content,
            }
        ],
        do_sample=False,
    )
    task_id = response.id
    task_status = ''
    get_cnt = 0

    # 轮询查询任务状态
    while task_status not in ['SUCCESS', 'FAILED'] and get_cnt <= 60:
        result_response = client.chat.asyncCompletions.retrieve_completion_result(id=task_id)
        task_status = result_response.task_status
        await asyncio.sleep(1)
        get_cnt += 1
    if task_status == 'SUCCESS':
        if return_json:
            _, json_data = try_parse_json_object(result_response.choices[0].message.content)
            return json_data
        else:
            return result_response.choices[0].message.content
    else:
        if return_json:
            return {}
        else:
            return ""
    

async def agenerate_all(client, user_contents, return_json=True):
    # 创建所有任务，并发执行
    tasks = [agenerate(client, user_content, return_json) for user_content in user_contents]
    all_json_data = await asyncio.gather(*tasks)
    return [json_data for json_data in all_json_data]


def generate_by_user_contents(client, user_contents, return_json=True, n_concurrencies=40):
    concurrent_contents = [user_contents[i:i+n_concurrencies] for i in range(0, len(user_contents), n_concurrencies)]
    all_results = []
    for user_contents in tqdm(concurrent_contents):
        results = asyncio.run(agenerate_all(client, user_contents, return_json))
        all_results += results
    return all_results


if __name__ == "__main__":
    from dotenv import load_dotenv
    load_dotenv()
    from zhipuai import ZhipuAI
    client = ZhipuAI()
    user_contents = [
        "作为童话之王，请以始终保持一颗善良的心为主题，写一篇简短的童话故事。故事应能激发孩子们的学习兴趣和想象力，同时帮助他们更好地理解和接受故事中蕴含的道德和价值观。",
        "请以‘科技改变生活’为主题写一篇短文。",
        "请以‘春天的旋律’为题写一首小诗。"
    ]
    #results = asyncio.run(agenerate_all(client, user_contents))
    results = generate_by_user_contents(client, user_contents, False, 40)
    print(results)